package com.ocr.paddle;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.ocr.paddle.domain.LocalHubOcrResultDTO;
import com.ocr.paddle.domain.OCRHubResultDTO;
import com.utils.Base64Util;
import com.utils.FileUtil;
import com.utils.HttpClientUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 文件名称： com.example.filecopy.LocalOcrUtils.java</br>
 * 初始作者： AirOrangeWorkSpace</br>
 * 创建日期： 2020/12/22 10:01</br>
 * 功能说明： TODO <br/>
 * =================================================<br/>
 * 修改记录：<br/>
 * 修改作者        日期       修改内容<br/>
 * ================================================<br/>
 * Copyright (c) 2020-2021 .All rights reserved.<br/>
 */
@Slf4j
public class LocalHubOcrUtils {

    public static List<LocalHubOcrResultDTO> localAllOcr(String textUrl, String faceUrl, List<File> allFiles) {
        long l = System.currentTimeMillis();
        List<LocalHubOcrResultDTO> localOcrResultDTOList = Lists.newArrayList();
        int urlCount = 0;
        boolean useTextUrl = false;
        boolean useFaceUrl = false;
        if (StringUtils.isEmpty(textUrl)) {
            log.info("不进行文字识别");
        } else {
            urlCount++;
            useTextUrl = true;
            log.info("需要进行文字识别");
        }
        if (StringUtils.isEmpty(faceUrl)) {
            log.info("不进行人脸识别");
        } else {
            urlCount++;
            useFaceUrl = true;
            log.info("需要进行人脸识别");
        }
        int ocrCount = allFiles.size() * urlCount;
        log.info("预估进行OCR" + ocrCount + "次");
        if (ocrCount == 0) {
            return localOcrResultDTOList;
        }

        int textIndex = 0;
        int faceIndex = 0;
        // 本地文件路径
        Map<String, String> partentFile = allFiles.stream().collect(Collectors.toMap(File::getName, File::getParent, (e1, e2) -> e1));
        for (File imageFile : allFiles) {
            List<OCRHubResultDTO> ocrHubResultDTOS = Lists.newArrayList();

            byte[] imgData = new byte[0];
            try {
                imgData = FileUtil.readFileByBytes(imageFile);
            } catch (IOException e) {
                log.error("图片读取错误");
                continue;
            }
            String imgStr = Base64Util.encode(imgData);
            JSONObject jsonObject = new JSONObject();
            List<String> imageParams = Lists.newArrayList();
            imageParams.add(imgStr);
            jsonObject.put("images", imageParams);

            String textResult = "";
            if (useTextUrl) {
                textResult = HttpClientUtils.sendJsonStr(textUrl, jsonObject.toJSONString());
                textIndex++;
                log.info("已进行文字OCR" + textIndex + "次");
                if (StringUtils.isEmpty(textResult)) {
                    log.error("获取文字接口失败");
                    log.info("还需进行" + (ocrCount - faceIndex - textIndex) + "次");
                    continue;
                }
                JSONObject resultTest = (JSONObject) JSONObject.parse(textResult);
                JSONArray textResultsArray = resultTest.getJSONArray("results");
                if (textResultsArray == null){
                    log.error("返回值错误，错误信息为：" + textResult);
                }

                for (Object o : textResultsArray) {
                    JSONObject j = (JSONObject) o;
                    JSONArray data = j.getJSONArray("data");
                    for (Object datum : data) {
                        JSONObject jo = (JSONObject) datum;
                        jo.remove("text_box_position");
                        OCRHubResultDTO ocrHubResultDTO = jo.toJavaObject(OCRHubResultDTO.class);
                        ocrHubResultDTOS.add(ocrHubResultDTO);
                    }
                }

            }
            int faceCount = 0;
            String faceResult = "";
            if (useFaceUrl) {
                faceResult = HttpClientUtils.sendJsonStr(faceUrl, jsonObject.toJSONString());
                faceIndex++;
                log.info("已进行人脸OCR" + faceIndex + "次");
                if (StringUtils.isEmpty(faceResult)) {
                    log.error("获取人脸识别接口失败");
                    log.info("还需进行" + (ocrCount - faceIndex - textIndex) + "次");
                    continue;
                }
                JSONObject resultFace = (JSONObject) JSONObject.parse(faceResult);
                JSONArray faceResultsArray = resultFace.getJSONArray("results");
                for (Object o : faceResultsArray) {
                    JSONObject j = (JSONObject) o;
                    JSONArray data = j.getJSONArray("data");
                    faceCount = data.size();
                }
            }

            log.info("还需进行" + (ocrCount - faceIndex - textIndex) + "次");
            localOcrResultDTOList.add(new LocalHubOcrResultDTO(imageFile.getAbsolutePath(), imageFile.getParent(), ocrHubResultDTOS, faceCount));
        }
        log.info("OCR 总耗时" + (System.currentTimeMillis() - l) / 1000 + "S");
        log.info("OCR 次数" + ocrCount + "次");
        log.info("OCR 平均耗时" + (System.currentTimeMillis() - l) / 1000 / ocrCount + "s");
        return localOcrResultDTOList;
    }

    private static void func(File file, List<File> fileList) {
        File[] fs = file.listFiles();
        for (File f : fs) {
            if (f.isDirectory()) {
                //若是目录，则递归打印该目录下的文件
                func(f, fileList);
            }
            if (f.isFile()) {
                String imageName = f.getName();
                boolean isJpg = imageName.endsWith(".jpg");
                boolean isJPG = imageName.endsWith(".JPG");
                boolean isPng = imageName.endsWith(".png");
                if (isJPG || isPng || isJpg) {
                    //若是图片加入列表
                    fileList.add(f);
                }
            }
        }
    }
}
